package com.wfh.mianshiji.ai.serializer;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.data.redis.serializer.RedisSerializer;
import org.springframework.data.redis.serializer.SerializationException;

import java.util.ArrayList;
import java.util.List;

/**
 * 使用 Kryo 的 Redis 序列化器（线程安全版本）
 * 通过 ThreadLocal 或池化 Kryo 实例保证线程安全
 * @author fenghuanwang
 */
public class KryoRedisSerializer<T> implements RedisSerializer<T> {

    private static final byte[] EMPTY_BYTE_ARRAY = new byte[0];

    // 使用 ThreadLocal 保证每个线程有自己的 Kryo 实例（Kryo 非线程安全）
    private static final ThreadLocal<Kryo> kryoLocal = ThreadLocal.withInitial(Kryo::new);

    // 可选：注册常用类以提升性能和兼容性（避免写类名）
    static {
        Kryo kryo = kryoLocal.get();
        kryo.register(MessageType.class);
        kryo.register(List.class);
        kryo.register(Message.class);
        kryo.register(ArrayList.class);
        kryo.register(UserMessage.class);
        kryo.register(AssistantMessage.class);
        // 示例：kryo.register(User.class);
        // 实际使用中建议在构造时传入注册类
    }

    private final Class<T> clazz;

    public KryoRedisSerializer(Class<T> clazz) {
        this.clazz = clazz;
    }

    @Override
    public byte[] serialize(T t) throws SerializationException {
        if (t == null) {
            return EMPTY_BYTE_ARRAY;
        }
        try (Output output = new Output(4096, -1)) {
            kryoLocal.get().writeObject(output, t);
            return output.toBytes();
        } catch (Exception e) {
            throw new SerializationException("Could not serialize: " + t, e);
        }
    }

    @Override
    public T deserialize(byte[] bytes) throws SerializationException {
        if (bytes == null || bytes.length == 0) {
            return null;
        }
        try (Input input = new Input(bytes)) {
            return kryoLocal.get().readObject(input, clazz);
        } catch (Exception e) {
            throw new SerializationException("Could not deserialize", e);
        }
    }
}